#include "ego_planner_adapter/trajectory_tracker.hpp"
#include <tf2/utils.h>
#include <algorithm>
#include <cmath>

// Include ego-planner messages if available
#ifdef TRAJ_UTILS_FOUND
#include "traj_utils/msg/bspline.hpp"
#endif

using namespace std::chrono_literals;

namespace ego_planner_adapter {

// PID Controller Implementation
PIDController::PIDController(double kp, double ki, double kd, double max_output, double min_output)
  : kp_(kp), ki_(ki), kd_(kd), max_output_(max_output), min_output_(min_output),
    integral_(0.0), previous_error_(0.0), first_call_(true) {}

double PIDController::update(double error, double dt) {
    if (dt <= 0.0) return 0.0;
    
    if (first_call_) {
        previous_error_ = error;
        first_call_ = false;
    }
    
    // Proportional term
    double proportional = kp_ * error;
    
    // Integral term with windup protection
    integral_ += error * dt;
    double max_integral = max_output_ / (ki_ + 1e-6);
    integral_ = std::clamp(integral_, -max_integral, max_integral);
    double integral_term = ki_ * integral_;
    
    // Derivative term
    double derivative = kd_ * (error - previous_error_) / dt;
    previous_error_ = error;
    
    // Total output
    double output = proportional + integral_term + derivative;
    return std::clamp(output, min_output_, max_output_);
}

void PIDController::reset() {
    integral_ = 0.0;
    previous_error_ = 0.0;
    first_call_ = true;
}

void PIDController::setGains(double kp, double ki, double kd) {
    kp_ = kp;
    ki_ = ki;
    kd_ = kd;
}

// TrajectoryTracker Implementation
TrajectoryTracker::TrajectoryTracker(const rclcpp::NodeOptions& options)
  : LifecycleNode("trajectory_tracker", options),
    current_position_(Eigen::Vector3d::Zero()),
    current_velocity_(Eigen::Vector3d::Zero()),
    current_orientation_(Eigen::Quaterniond::Identity()),
    current_yaw_(0.0),
    has_odom_(false),
    has_trajectory_(false) {
    
    RCLCPP_INFO(get_logger(), "TrajectoryTracker node created");
}

rclcpp_lifecycle::node_interfaces::LifecycleNodeInterface::CallbackReturn
TrajectoryTracker::on_configure(const rclcpp_lifecycle::State& /*state*/) {
    RCLCPP_INFO(get_logger(), "Configuring TrajectoryTracker");
    
    // Declare parameters
    declare_parameter("control_frequency", 50.0);
    declare_parameter("max_linear_velocity", 2.0);
    declare_parameter("max_angular_velocity", 1.0);
    declare_parameter("position_tolerance", 0.1);
    declare_parameter("yaw_tolerance", 0.1);
    declare_parameter("trajectory_timeout", 5.0);
    declare_parameter("odom_frame", "odom");
    declare_parameter("base_frame", "base_footprint");
    
    // PID parameters
    declare_parameter("pid_x.kp", 2.0);
    declare_parameter("pid_x.ki", 0.1);
    declare_parameter("pid_x.kd", 0.5);
    declare_parameter("pid_y.kp", 2.0);
    declare_parameter("pid_y.ki", 0.1);
    declare_parameter("pid_y.kd", 0.5);
    declare_parameter("pid_z.kp", 3.0);
    declare_parameter("pid_z.ki", 0.1);
    declare_parameter("pid_z.kd", 0.8);
    declare_parameter("pid_yaw.kp", 2.0);
    declare_parameter("pid_yaw.ki", 0.0);
    declare_parameter("pid_yaw.kd", 0.3);
    
    // Get parameters
    control_frequency_ = get_parameter("control_frequency").as_double();
    max_linear_velocity_ = get_parameter("max_linear_velocity").as_double();
    max_angular_velocity_ = get_parameter("max_angular_velocity").as_double();
    position_tolerance_ = get_parameter("position_tolerance").as_double();
    yaw_tolerance_ = get_parameter("yaw_tolerance").as_double();
    trajectory_timeout_ = get_parameter("trajectory_timeout").as_double();
    odom_frame_ = get_parameter("odom_frame").as_string();
    base_frame_ = get_parameter("base_frame").as_string();
    
    // Initialize PID controllers
    double kp_x = get_parameter("pid_x.kp").as_double();
    double ki_x = get_parameter("pid_x.ki").as_double();
    double kd_x = get_parameter("pid_x.kd").as_double();
    
    double kp_y = get_parameter("pid_y.kp").as_double();
    double ki_y = get_parameter("pid_y.ki").as_double();
    double kd_y = get_parameter("pid_y.kd").as_double();
    
    double kp_z = get_parameter("pid_z.kp").as_double();
    double ki_z = get_parameter("pid_z.ki").as_double();
    double kd_z = get_parameter("pid_z.kd").as_double();
    
    double kp_yaw = get_parameter("pid_yaw.kp").as_double();
    double ki_yaw = get_parameter("pid_yaw.ki").as_double();
    double kd_yaw = get_parameter("pid_yaw.kd").as_double();
    
    pid_x_ = std::make_unique<PIDController>(kp_x, ki_x, kd_x, max_linear_velocity_, -max_linear_velocity_);
    pid_y_ = std::make_unique<PIDController>(kp_y, ki_y, kd_y, max_linear_velocity_, -max_linear_velocity_);
    pid_z_ = std::make_unique<PIDController>(kp_z, ki_z, kd_z, max_linear_velocity_, -max_linear_velocity_);
    pid_yaw_ = std::make_unique<PIDController>(kp_yaw, ki_yaw, kd_yaw, max_angular_velocity_, -max_angular_velocity_);
    
    // Initialize trajectory interpolator
    trajectory_interpolator_ = std::make_unique<BSplineInterpolator>();
    
    // Create subscribers
    odom_sub_ = create_subscription<nav_msgs::msg::Odometry>(
        "odometry", 10,
        std::bind(&TrajectoryTracker::odometryCallback, this, std::placeholders::_1)
    );
    
#ifdef TRAJ_UTILS_FOUND
    trajectory_sub_ = create_subscription<traj_utils::msg::Bspline>(
        "trajectory", 10,
        std::bind(&TrajectoryTracker::trajectoryCallback, this, std::placeholders::_1)
    );
#else
    RCLCPP_WARN(get_logger(), "traj_utils not found, trajectory subscription disabled");
#endif
    
    // Create publishers
    cmd_vel_pub_ = create_publisher<geometry_msgs::msg::Twist>("cmd_vel", 10);
    target_altitude_pub_ = create_publisher<std_msgs::msg::Float64>("target_altitude", 10);
    viz_pub_ = create_publisher<visualization_msgs::msg::Marker>("trajectory_marker", 10);
    
    RCLCPP_INFO(get_logger(), "TrajectoryTracker configured successfully");
    return rclcpp_lifecycle::node_interfaces::LifecycleNodeInterface::CallbackReturn::SUCCESS;
}

rclcpp_lifecycle::node_interfaces::LifecycleNodeInterface::CallbackReturn
TrajectoryTracker::on_activate(const rclcpp_lifecycle::State& /*state*/) {
    RCLCPP_INFO(get_logger(), "Activating TrajectoryTracker");
    
    // Activate publishers
    cmd_vel_pub_->on_activate();
    target_altitude_pub_->on_activate();
    viz_pub_->on_activate();
    
    // Create control timer
    auto timer_period = std::chrono::duration<double>(1.0 / control_frequency_);
    control_timer_ = create_wall_timer(
        std::chrono::duration_cast<std::chrono::nanoseconds>(timer_period),
        std::bind(&TrajectoryTracker::controlTimerCallback, this)
    );
    
    RCLCPP_INFO(get_logger(), "TrajectoryTracker activated successfully");
    return rclcpp_lifecycle::node_interfaces::LifecycleNodeInterface::CallbackReturn::SUCCESS;
}

rclcpp_lifecycle::node_interfaces::LifecycleNodeInterface::CallbackReturn
TrajectoryTracker::on_deactivate(const rclcpp_lifecycle::State& /*state*/) {
    RCLCPP_INFO(get_logger(), "Deactivating TrajectoryTracker");
    
    // Stop timer
    control_timer_.reset();
    
    // Deactivate publishers
    cmd_vel_pub_->on_deactivate();
    target_altitude_pub_->on_deactivate();
    viz_pub_->on_deactivate();
    
    // Emergency stop
    emergencyStop();
    
    RCLCPP_INFO(get_logger(), "TrajectoryTracker deactivated");
    return rclcpp_lifecycle::node_interfaces::LifecycleNodeInterface::CallbackReturn::SUCCESS;
}

#ifdef TRAJ_UTILS_FOUND
void TrajectoryTracker::trajectoryCallback(const std::shared_ptr<const traj_utils::msg::Bspline>& msg) {
    std::lock_guard<std::mutex> lock(trajectory_mutex_);
    
    RCLCPP_INFO(get_logger(), "Received new trajectory with %zu control points", msg->pos_pts.size());
    
    // Set new trajectory
    if (trajectory_interpolator_->setTrajectory(msg)) {
        has_trajectory_ = true;
        trajectory_start_time_ = now();
        last_trajectory_time_ = now();
        
        // Reset PID controllers for new trajectory
        resetControllers();
        
        RCLCPP_INFO(get_logger(), "Trajectory set successfully, duration: %.2f seconds",
                   trajectory_interpolator_->getEndTime() - trajectory_interpolator_->getStartTime());
    } else {
        RCLCPP_WARN(get_logger(), "Failed to set trajectory");
        has_trajectory_ = false;
    }
}
#else
void TrajectoryTracker::trajectoryCallback(const std::shared_ptr<const geometry_msgs::msg::PoseStamped>& /*msg*/) {
    RCLCPP_WARN(get_logger(), "Trajectory callback called but traj_utils not available");
}
#endif

void TrajectoryTracker::odometryCallback(const std::shared_ptr<const nav_msgs::msg::Odometry>& msg) {
    std::lock_guard<std::mutex> lock(state_mutex_);
    
    // Update current state
    current_position_.x() = msg->pose.pose.position.x;
    current_position_.y() = msg->pose.pose.position.y;
    current_position_.z() = msg->pose.pose.position.z;
    
    current_velocity_.x() = msg->twist.twist.linear.x;
    current_velocity_.y() = msg->twist.twist.linear.y;
    current_velocity_.z() = msg->twist.twist.linear.z;
    
    current_orientation_.x() = msg->pose.pose.orientation.x;
    current_orientation_.y() = msg->pose.pose.orientation.y;
    current_orientation_.z() = msg->pose.pose.orientation.z;
    current_orientation_.w() = msg->pose.pose.orientation.w;
    
    // Extract yaw from quaternion
    current_yaw_ = tf2::getYaw(msg->pose.pose.orientation);
    
    has_odom_ = true;
    last_trajectory_time_ = now();
}

void TrajectoryTracker::controlTimerCallback() {
    if (get_current_state().id() == lifecycle_msgs::msg::State::PRIMARY_STATE_ACTIVE) {
        updateControl();
    }
}

void TrajectoryTracker::updateControl() {
    std::lock_guard<std::mutex> state_lock(state_mutex_);
    std::lock_guard<std::mutex> traj_lock(trajectory_mutex_);
    
    // Check prerequisites
    if (!has_odom_) {
        RCLCPP_WARN_THROTTLE(get_logger(), *get_clock(), 1000, "No odometry data received");
        return;
    }
    
    if (!has_trajectory_) {
        // Publish zero velocity
        publishVelocityCommand(Eigen::Vector3d::Zero(), 0.0);
        return;
    }
    
    if (!isTrajectoryValid()) {
        RCLCPP_WARN_THROTTLE(get_logger(), *get_clock(), 1000, "Trajectory expired or invalid");
        emergencyStop();
        return;
    }
    
    // Get current time relative to trajectory start
    double current_time = getCurrentTime();
    
    // Get target position and velocity from trajectory
    Eigen::Vector3d target_position, target_velocity;
    double target_yaw;
    
    if (!trajectory_interpolator_->getPositionAtTime(current_time, target_position) ||
        !trajectory_interpolator_->getVelocityAtTime(current_time, target_velocity) ||
        !trajectory_interpolator_->getYawAtTime(current_time, target_yaw)) {
        
        RCLCPP_WARN(get_logger(), "Failed to interpolate trajectory at time %.3f", current_time);
        emergencyStop();
        return;
    }
    
    // Compute position errors
    Eigen::Vector3d position_error = target_position - current_position_;
    double yaw_error = target_yaw - current_yaw_;
    
    // Normalize yaw error to [-pi, pi]
    while (yaw_error > M_PI) yaw_error -= 2 * M_PI;
    while (yaw_error < -M_PI) yaw_error += 2 * M_PI;
    
    // Compute control commands
    double dt = 1.0 / control_frequency_;
    
    Eigen::Vector3d control_velocity;
    control_velocity.x() = pid_x_->update(position_error.x(), dt) + target_velocity.x();
    control_velocity.y() = pid_y_->update(position_error.y(), dt) + target_velocity.y();
    control_velocity.z() = pid_z_->update(position_error.z(), dt) + target_velocity.z();
    
    double control_yaw = pid_yaw_->update(yaw_error, dt);
    
    // Apply velocity limits
    double linear_speed = control_velocity.norm();
    if (linear_speed > max_linear_velocity_) {
        control_velocity *= (max_linear_velocity_ / linear_speed);
    }
    control_yaw = std::clamp(control_yaw, -max_angular_velocity_, max_angular_velocity_);
    
    // Publish control commands
    publishVelocityCommand(control_velocity, control_yaw);
    
    // Publish target altitude
    auto altitude_msg = std::make_unique<std_msgs::msg::Float64>();
    altitude_msg->data = target_position.z();
    target_altitude_pub_->publish(std::move(altitude_msg));
    
    // Publish visualization
    publishVisualization();
    
    // Debug output
    RCLCPP_DEBUG(get_logger(), 
                "Control: pos_err=[%.3f,%.3f,%.3f], yaw_err=%.3f, cmd_vel=[%.3f,%.3f,%.3f], cmd_yaw=%.3f",
                position_error.x(), position_error.y(), position_error.z(), yaw_error,
                control_velocity.x(), control_velocity.y(), control_velocity.z(), control_yaw);
}

void TrajectoryTracker::publishVelocityCommand(const Eigen::Vector3d& linear_vel, double angular_vel_z) {
    auto cmd_msg = std::make_unique<geometry_msgs::msg::Twist>();
    cmd_msg->linear.x = linear_vel.x();
    cmd_msg->linear.y = linear_vel.y();
    cmd_msg->linear.z = linear_vel.z();
    cmd_msg->angular.z = angular_vel_z;
    
    cmd_vel_pub_->publish(std::move(cmd_msg));
}

void TrajectoryTracker::publishVisualization() {
    if (!has_trajectory_ || !trajectory_interpolator_->isValid()) return;
    
    auto marker = std::make_unique<visualization_msgs::msg::Marker>();
    marker->header.frame_id = odom_frame_;
    marker->header.stamp = now();
    marker->ns = "trajectory";
    marker->id = 0;
    marker->type = visualization_msgs::msg::Marker::LINE_STRIP;
    marker->action = visualization_msgs::msg::Marker::ADD;
    
    // Set marker properties
    marker->scale.x = 0.05;
    marker->color.r = 1.0;
    marker->color.g = 0.0;
    marker->color.b = 0.0;
    marker->color.a = 1.0;
    
    // Add trajectory points
    double start_time = trajectory_interpolator_->getStartTime();
    double end_time = trajectory_interpolator_->getEndTime();
    double dt = 0.1; // 10 Hz visualization
    
    for (double t = start_time; t <= end_time; t += dt) {
        Eigen::Vector3d pos;
        if (trajectory_interpolator_->getPositionAtTime(t, pos)) {
            geometry_msgs::msg::Point point;
            point.x = pos.x();
            point.y = pos.y();
            point.z = pos.z();
            marker->points.push_back(point);
        }
    }
    
    viz_pub_->publish(std::move(marker));
}

void TrajectoryTracker::emergencyStop() {
    publishVelocityCommand(Eigen::Vector3d::Zero(), 0.0);
    RCLCPP_WARN(get_logger(), "Emergency stop executed");
}

void TrajectoryTracker::resetControllers() {
    pid_x_->reset();
    pid_y_->reset();
    pid_z_->reset();
    pid_yaw_->reset();
}

bool TrajectoryTracker::isTrajectoryValid() const {
    if (!has_trajectory_ || !trajectory_interpolator_->isValid()) return false;
    
    // Check timeout
    auto time_since_trajectory = (now() - trajectory_start_time_).seconds();
    if (time_since_trajectory > trajectory_timeout_) return false;
    
    return true;
}

double TrajectoryTracker::getCurrentTime() const {
    return (now() - trajectory_start_time_).seconds();
}

} // namespace ego_planner_adapter